# Import required modules
import cv2
import numpy as np
import os
import glob

# Define the dimensions of checkerboard
CHECKERBOARD = (11, 8)

# stop the iteration when specified
# accuracy, epsilon, is reached or
# specified number of iterations are completed.
criteria = (cv2.TERM_CRITERIA_EPS +
            cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)

# Vector for 3D points
threedpoints = []

# Vector for 2D points
twodpoints = []

#  3D points real world coordinates
objectp3d = np.zeros((1, CHECKERBOARD[0]
                      * CHECKERBOARD[1],
                      3), np.float32)
objectp3d[0, :, :2] = np.mgrid[0:CHECKERBOARD[0],
                      0:CHECKERBOARD[1]].T.reshape(-1, 2)
prev_img_shape = None

# Extracting path of individual image stored
# in a given directory. Since no path is
# specified, it will take current directory
# jpg files alone

def calibrate_camera(images):
    for filename in images:
        print(filename)
        image = cv2.imread(filename)
        grayColor = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

        # Find the chess board corners
        # If desired number of corners are
        # found in the image then ret = true
        ret, corners = cv2.findChessboardCorners(
            grayColor, CHECKERBOARD,
            cv2.CALIB_CB_ADAPTIVE_THRESH
            + cv2.CALIB_CB_FAST_CHECK +
            cv2.CALIB_CB_NORMALIZE_IMAGE)

        # If desired number of corners can be detected then,
        # refine the pixel coordinates and display
        # them on the images01 of checker board
        if ret == True:
            threedpoints.append(objectp3d)

            # Refining pixel coordinates
            # for given 2d points.
            corners2 = cv2.cornerSubPix(
                grayColor, corners, (11, 11), (-1, -1), criteria)

            twodpoints.append(corners2)

            # Draw and display the corners
            image = cv2.drawChessboardCorners(image,
                                              CHECKERBOARD,
                                              corners2, ret)

        cv2.imshow('img', image)
        k = cv2.waitKey(0)
        if k & 0xFF == ord('q'):
            cv2.destroyAllWindows()
            break
    cv2.destroyAllWindows()

    h, w = cv2.imread(images[0]).shape[:2]


    # Perform camera calibration by
    # passing the value of above found out 3D points (threedpoints)
    # and its corresponding pixel coordinates of the
    # detected corners (twodpoints)
    ret, matrix, distortion, r_vecs, t_vecs = cv2.calibrateCamera(
        threedpoints, twodpoints, grayColor.shape[::-1], None, None)

    print("ret:", ret)
    # Displaying required output
    print(" Camera matrix:")
    print(matrix)

    print("\n Distortion coefficient:")
    print(distortion)

    print("\n Rotation Vectors:")
    print(r_vecs)

    print("\n Translation Vectors:")
    print(t_vecs)

    print("FOV:", fov(matrix))
    return {
        "matrix": matrix,
        "distortion": distortion,
        "r_vecs": r_vecs,
        "t_vecs": t_vecs
    }

def fov(matrix):
    import math
    fx = matrix[0, 0]
    fy = matrix[1, 1]
    w = 1280
    h = 720
    fov_x = 2 * math.atan(w / (2 * fx)) * 180 / math.pi
    fov_y = 2 * math.atan(h / (2 * fy)) * 180 / math.pi
    fov_d = math.sqrt(fov_x * fov_x + fov_y * fov_y)
    return fov_x, fov_y, fov_d

if __name__ == "__main__":
    import pickle
    images = glob.glob('images_tmp/*.jpg')
    result = calibrate_camera(images)
    pickle.dump(result, open('camera_tmp.pkl', 'wb'))

    # images = glob.glob('data/camera/left*.png')
    # result = calibrate_camera(images)
    # pickle.dump(result, open('data/camera/calibration/left_origin.pkl', 'wb'))

    # IMX 385
    #  Camera matrix:
# [[704.19697079   0.         918.34752338]
#  [  0.         704.64503299 563.32098459]
#  [  0.           0.           1.        ]]

#  Distortion coefficient:
# [[-8.23821460e-03  2.75231331e-03 -1.70803300e-04  2.90067585e-05
#   -1.05421462e-02]]
    # FOV: (107.47703664932432, 74.9290177820396)
    # 

# second: 
# FOV: (89.43559582810146, 58.23558208066522)
#  Distortion coefficient:
# [[ 1.75113569e-01 -2.50387154e-01  2.44603016e-04  7.73205247e-04
#    1.15175502e-01]]
# 
# Third
# FOV: (80.68952993042073, 51.07980247627678)
#  Camera matrix:
# [[1.13020188e+03 0.00000000e+00 9.66002715e+02]
#  [0.00000000e+00 1.13010747e+03 5.51514951e+02]
#  [0.00000000e+00 0.00000000e+00 1.00000000e+00]]
#  Distortion coefficient:
# [[ 1.57336457e-01 -2.99111846e-01  1.34622563e-04  8.39214351e-04
#    2.72836595e-01]]

# FOV: (80.67845608225454, 51.04823422171724)
# Distortion coefficient:
# [[ 1.33875855e-01 -1.90566436e-01 -1.85070386e-04  1.25977088e-03
#    8.57243190e-02]]
#  Camera matrix:
# [[1.13042326e+03 0.00000000e+00 9.66652533e+02]
#  [0.00000000e+00 1.13090824e+03 5.31836116e+02]
#  [0.00000000e+00 0.00000000e+00 1.00000000e+00]]

# FPS 240
# FOV: (86.34856386263988, 55.618679691646705, 102.0280392356678)
#  Camera matrix:
# [[682.14411131   0.         617.78650712]
#  [  0.         682.53104944 350.13208729]
#  [  0.           0.           1.        ]]

# FPS 120
# FOV: (67.4451317433, 41.16271978740408, 79.0140196178339)
#  Camera matrix:
# [[1.43823183e+03 0.00000000e+00 9.52357948e+02]
#  [0.00000000e+00 1.43806709e+03 5.78416269e+02]
#  [0.00000000e+00 0.00000000e+00 1.00000000e+00]]

#  Distortion coefficient:
# [[ 8.14450618e-02 -8.80969415e-02 -4.48229718e-04 -4.21658823e-05
#   -1.06300847e-01]]